package com.intelligent.security;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;

/**
 * Minimal JWT HS256 utility to create/verify tokens and extract userId.
 * Token format: header.payload.signature where payload contains: sub (userId),
 * exp (ms epoch)
 */
public class JwtUtil {
    private static final String HMAC_ALGO = "HmacSHA256";
    // IMPORTANT: change this secret in production. Can be overridden by env var
    // JWT_SECRET
    private static final String SECRET = System.getenv().getOrDefault("JWT_SECRET", "change_this_secret_please");
    // default 7 days
    private static final long EXPIRATION_MS = 7L * 24 * 60 * 60 * 1000;

    public static String createToken(Long userId) {
        long exp = System.currentTimeMillis() + EXPIRATION_MS;
        Map<String, Object> payload = new HashMap<>();
        payload.put("sub", String.valueOf(userId));
        payload.put("exp", exp);
        return createToken(payload);
    }

    public static String createToken(Long userId, String userAccount) {
        long exp = System.currentTimeMillis() + EXPIRATION_MS;
        Map<String, Object> payload = new HashMap<>();
        payload.put("sub", String.valueOf(userId));
        payload.put("account", userAccount);
        payload.put("exp", exp);
        return createToken(payload);
    }

    public static boolean validate(String token) {
        try {
            Map<String, Object> payload = parseAndVerify(token);
            if (payload == null)
                return false;
            Object expObj = payload.get("exp");
            if (expObj instanceof Number) {
                long exp = ((Number) expObj).longValue();
                return System.currentTimeMillis() < exp;
            }
            if (expObj instanceof String) {
                long exp = Long.parseLong((String) expObj);
                return System.currentTimeMillis() < exp;
            }
            return false;
        } catch (Exception e) {
            return false;
        }
    }

    public static Long getUserId(String token) {
        try {
            Map<String, Object> payload = parseAndVerify(token);
            if (payload == null)
                return null;
            Object sub = payload.get("sub");
            if (sub == null)
                return null;
            return Long.parseLong(String.valueOf(sub));
        } catch (Exception e) {
            return null;
        }
    }

    public static String getUserAccount(String token) {
        try {
            Map<String, Object> payload = parseAndVerify(token);
            if (payload == null)
                return null;
            Object account = payload.get("account");
            return account != null ? String.valueOf(account) : null;
        } catch (Exception e) {
            return null;
        }
    }

    private static String createToken(Map<String, Object> payload) {
        String headerJson = "{\"alg\":\"HS256\",\"typ\":\"JWT\"}";
        String payloadJson = toJson(payload);
        String headerB64 = base64UrlEncode(headerJson.getBytes(StandardCharsets.UTF_8));
        String payloadB64 = base64UrlEncode(payloadJson.getBytes(StandardCharsets.UTF_8));
        String signingInput = headerB64 + "." + payloadB64;
        String signature = sign(signingInput);
        return signingInput + "." + signature;
    }

    private static Map<String, Object> parseAndVerify(String token) {
        String[] parts = token == null ? new String[0] : token.split("\\.");
        if (parts.length != 3)
            return null;
        String signingInput = parts[0] + "." + parts[1];
        String expectedSig = sign(signingInput);
        if (!constantTimeEquals(expectedSig, parts[2]))
            return null;
        String payloadJson = new String(base64UrlDecode(parts[1]), StandardCharsets.UTF_8);
        return fromJson(payloadJson);
    }

    private static String sign(String data) {
        try {
            Mac mac = Mac.getInstance(HMAC_ALGO);
            mac.init(new SecretKeySpec(SECRET.getBytes(StandardCharsets.UTF_8), HMAC_ALGO));
            byte[] raw = mac.doFinal(data.getBytes(StandardCharsets.UTF_8));
            return base64UrlEncode(raw);
        } catch (Exception e) {
            throw new RuntimeException("Failed to sign JWT", e);
        }
    }

    private static boolean constantTimeEquals(String a, String b) {
        if (a == null || b == null)
            return false;
        if (a.length() != b.length())
            return false;
        int result = 0;
        for (int i = 0; i < a.length(); i++) {
            result |= a.charAt(i) ^ b.charAt(i);
        }
        return result == 0;
    }

    private static String base64UrlEncode(byte[] bytes) {
        return Base64.getUrlEncoder().withoutPadding().encodeToString(bytes);
    }

    private static byte[] base64UrlDecode(String str) {
        return Base64.getUrlDecoder().decode(str);
    }

    // Very small JSON helpers to avoid dependencies. Only handles flat maps with
    // String/Number values.
    private static String toJson(Map<String, Object> map) {
        StringBuilder sb = new StringBuilder();
        sb.append('{');
        boolean first = true;
        for (Map.Entry<String, Object> e : map.entrySet()) {
            if (!first)
                sb.append(',');
            first = false;
            sb.append('"').append(escape(e.getKey())).append('"').append(':');
            Object v = e.getValue();
            if (v instanceof Number || v instanceof Boolean) {
                sb.append(String.valueOf(v));
            } else {
                sb.append('"').append(escape(String.valueOf(v))).append('"');
            }
        }
        sb.append('}');
        return sb.toString();
    }

    private static Map<String, Object> fromJson(String json) {
        Map<String, Object> map = new HashMap<>();
        String trimmed = json.trim();
        if (trimmed.length() < 2 || trimmed.charAt(0) != '{' || trimmed.charAt(trimmed.length() - 1) != '}') {
            return map;
        }
        String body = trimmed.substring(1, trimmed.length() - 1).trim();
        if (body.isEmpty())
            return map;
        String[] parts = body.split(",");
        for (String part : parts) {
            String[] kv = part.split(":", 2);
            if (kv.length != 2)
                continue;
            String key = unquote(kv[0].trim());
            String val = kv[1].trim();
            if (val.startsWith("\"") && val.endsWith("\"")) {
                map.put(key, unquote(val));
            } else {
                try {
                    if (val.contains(".")) {
                        map.put(key, Double.parseDouble(val));
                    } else {
                        map.put(key, Long.parseLong(val));
                    }
                } catch (NumberFormatException e) {
                    map.put(key, val);
                }
            }
        }
        return map;
    }

    private static String unquote(String s) {
        String t = s;
        if (t.startsWith("\"") && t.endsWith("\"")) {
            t = t.substring(1, t.length() - 1);
        }
        return t.replace("\\\"", "\"").replace("\\\\", "\\");
    }

    private static String escape(String s) {
        return s.replace("\\", "\\\\").replace("\"", "\\\"");
    }
}
